#!/usr/bin/env python3
# models/generator.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from losses.custom_losses import RRDB   # 导入RRDB类，避免重复定义

class RRDBNet(nn.Module):
    """ESRGAN 生成器的基础模块（RRDB 网络）"""
    def __init__(self, num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4):
        super(RRDBNet, self).__init__()
        self.scale = scale
        # 示例结构：卷积 + RRDB块 + 上采样 + 输出卷积
        self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1, bias=True)
        self.body = self._make_rrdb_blocks(num_feat, num_block, num_grow_ch)
        self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=True)
        self.upsampler = self._make_upsampler(num_feat, scale)
        self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1, bias=True)

    def _make_rrdb_blocks(self, num_feat, num_block, num_grow_ch):
        blocks = []
        for _ in range(num_block):
            blocks.append(RRDB(num_feat, num_grow_ch))
        return nn.Sequential(*blocks)

    def _make_upsampler(self, num_feat, scale):
        # 实现上采样模块（如PixelShuffle）
        upsampler = []
        for _ in range(int(torch.log2(torch.tensor(scale)))):
            upsampler.append(nn.Conv2d(num_feat, num_feat * 4, 3, 1, 1, bias=True))
            upsampler.append(nn.PixelShuffle(2))
        return nn.Sequential(*upsampler)

    def forward(self, x):
        # 实现前向传播逻辑
        feat = self.conv_first(x)
        body_feat = self.conv_body(self.body(feat))
        feat = feat + body_feat
        out = self.conv_last(self.upsampler(feat))
        return out

# 定义ESRGAN生成器（继承RRDB网络，保持接口一致性）
class ESRGAN(RRDBNet):
    """ESRGAN生成器类（与RRDB网络结构一致，用于统一接口）"""
    # 在 ESRGAN 类中补充
    def __init__(self, scale_factor=4, num_block=23, num_grow_ch=32, **kwargs):
        super(ESRGAN, self).__init__(
            scale=scale_factor,
            num_block=num_block,
            num_grow_ch=num_grow_ch,** kwargs
        )
        self.scale_factor = scale_factor
        self.conv_first = nn.Conv2d(3, 64, 3, 1, 1, bias=True)
        # 保存参数为实例变量，供后续调用
        self.num_rrdb_blocks = num_block  # RRDB块数量
        self.num_grow_ch = num_grow_ch    # 增长通道数
        # 正确调用_make_rrdb_blocks，使用实例变量
        self.RRDB_trunk = self._make_rrdb_blocks(64, self.num_rrdb_blocks, self.num_grow_ch)
        self.trunk_conv = nn.Conv2d(64, 64, 3, 1, 1, bias=True)
        self.HRconv = nn.Conv2d(64, 64, 3, 1, 1, bias=True)
        self.conv_last = nn.Conv2d(64, 3, 3, 1, 1, bias=True)
        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
        
        # 新增：根据缩放因子添加上采样层
        self.upsampler = self._make_upsampler(64, scale_factor)     # 复用父类的上采样方法

    def forward(self, x):
        fea = self.conv_first(x)
        trunk = self.trunk_conv(self.RRDB_trunk(fea))
        fea = fea + trunk
        # 上采样逻辑
        """
        @date 2025/09/29
        @author cmx-cxd
        添加上采样步骤
        """
        fea = self.upsampler(fea)           # 先上采样到高分辨率尺寸
        
        fea = self.lrelu(self.HRconv(fea))
        out = self.conv_last(fea)
        return out